import torch
import numpy as np
from vae import VAE
from torchvision.transforms import transforms
import cv2

def prediction(model:VAE, img_path, save_path, device):
    model.eval()
    img = cv2.imread(img_path)
    img = cv2.resize(img, [128, 128])
    h, w = img.shape[:2]

    
    inp = transforms.ToTensor()(img)
    inp = inp.unsqueeze(0)
    
    with torch.no_grad():
        out = model.pred(inp.to(device))

    out = torch.squeeze(out)
    img2 = out.to('cpu')
    img2 = img2.numpy()
    img2 = (img2 * 255).astype(np.uint8)
    img2 = np.transpose(img2, (1, 2, 0))
    # img2 = cv2.cvtColor(img2, cv2.COLOR_RGB2BGR)

    # 拼接输出和输入
    h, w = img.shape[:2]
    h2, w2 = img2.shape[:2]
    m1 = np.zeros((max(h, h2), max(w, w2), 3), dtype=np.uint8)
    m2 = np.zeros((max(h, h2), max(w, w2), 3), dtype=np.uint8)
    m1[:h, :w, :] = img
    m2[:h2, :w2, :] = img2
    cat_img = img = np.concatenate([m1, m2], axis=1)

    cv2.imwrite(save_path, cat_img)

